# coding=utf-8
import bytedance.context
from euler import base_compat_middleware, client

from apitest.engine.client import Client
from apitest.engine.env import run_env
from apitest.engine.middleware import rpc_log_middleware


class RpcClient(Client):

    def __init__(self, service_cls, psm, cluster="default", timeout=10, **kwargs):
        super(RpcClient, self).__init__()

        bytedance.context.set("env", run_env.env_tag)  # 设置环境
        kwargs["timeout"] = timeout
        self.service_cls = service_cls
        self.psm = psm
        self.cluster = cluster
        self.target = "sd://{}?cluster={}".format(self.psm, self.cluster)
        self.client = client.Client(self.service_cls, self.target, disable_logging=False, **kwargs)
        self.client.use(base_compat_middleware.client_middleware)
        # self.use(rpc_response_error_middleware)
        self.client.use(rpc_log_middleware)
